package task

import (
	"fmt"
	"log"
	"net"
	"port-scanner/util"
	"port-scanner/vars"
	"sync"
	"time"
)

type Task struct {
	IP   string
	Port int
}

func (*Task) Connect(ip string, port int) (net.Conn, error) {
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), 5*time.Second)

	defer func() {
		if conn != nil {
			conn.Close()
		}
	}()

	return conn, err
}

func MakeTaskGroups(ipList []net.IP, ports []int) [][]Task {
	tasks := make([]Task, 0, len(ipList)*len(ports))
	for _, ip := range ipList {
		for _, port := range ports {
			tasks = append(tasks, Task{ip.String(), port})
		}
	}
	taskGroups := util.BatchSplit(tasks, vars.TaskBatchSize)
	return taskGroups
}

func RunTaskGroup(group []Task) {
	var wg sync.WaitGroup
	wg.Add(len(group))
	for _, task := range group {
		go func(task Task) {
			defer wg.Done()
			_, err := task.Connect(task.IP, task.Port)
			if err == nil {
				log.Printf("ip: %s port: %d is open\n", task.IP, task.Port)
				if v, ok := vars.Result.Load(task.IP); ok {
					if ports, ok1 := v.([]int); ok1 {
						ports = append(ports, task.Port)
						vars.Result.Store(task.IP, ports)
					}
				} else {
					ports := []int{task.Port}
					vars.Result.Store(task.IP, ports)
				}
			}
		}(task)
	}
	wg.Wait()
}
